import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
from affine import Affine
from osgeo import gdal, osr
from rasterio import MemoryFile
from rasterio.mask import mask
from shapely.geometry import mapping


def show_data(data, cmap=plt.cm.jet):
    """
    显示一个 numpy 数组并允许指定 cmap
    :param data: 要显示的 numpy 数组
    :param cmap: 色彩映射，默认为 plt.cm.jet
    """
    plt.imshow(data, cmap=cmap)
    plt.colorbar()
    plt.show()


PNG_DPI = 300


def export_png(data, png_path, width=None, height=None, full_image=False, cmap=plt.cm.jet):
    """
    将 numpy 数组导出为 PNG 图片
    :param data: 要导出的 numpy 数组
    :param png_path: 保存 PNG 图片的路径
    :param width: 图片宽度（像素）
    :param height: 图片高度（像素）
    :param full_image: 是否只保存图片而不包含颜色条等（默认 False）
    :param cmap: 色彩映射，默认为 plt.cm.jet
    """
    fig, ax = plt.subplots()
    if width and height:
        fig.set_size_inches(width / PNG_DPI, height / PNG_DPI)  # 以像素为单位指定宽高
    img = ax.imshow(data, cmap=cmap, origin='upper')

    if full_image:
        ax.axis('off')  # 关闭坐标轴
        plt.savefig(png_path, bbox_inches='tight', pad_inches=0, dpi=PNG_DPI)
    else:
        fig.colorbar(img, ax=ax)
        plt.savefig(png_path, bbox_inches='tight', pad_inches=0.1, dpi=PNG_DPI)

    plt.close()


def geo_transform_to_affine(geo_transform):
    """
    将 gdal 的 geo_transform 格式转换为 rasterio 的 Affine 变换格式
    :param geo_transform: gdal 的 geo_transform 格式
    :return: rasterio 的 Affine 变换格式
    """
    return Affine(geo_transform[1], geo_transform[2], geo_transform[0],
                  geo_transform[4], geo_transform[5], geo_transform[3])


def affine_to_geo_transform(affine):
    """
    将 rasterio 的 Affine 变换转换为 gdal 的 geo_transform 格式
    :param affine: rasterio 的 Affine 变换
    :return: gdal 的 geo_transform 格式
    """
    return affine.c, affine.a, affine.b, affine.f, affine.d, affine.e


def apply_shapefile_mask(data, affine, shapefile_path, nodata_value=np.nan):
    """
    使用 shapefile 对 numpy 数组进行蒙版处理
    :param data: 要进行蒙版处理的 numpy 数组
    :param affine: 仿射变换矩阵
    :param shapefile_path: shapefile 文件路径
    :param nodata_value: 蒙版区域的填充值，默认为 np.nan
    :return: 处理后的 numpy 数组和新的仿射变换矩阵
    """
    # 读取 shapefile 文件
    shapefile = gpd.read_file(shapefile_path)
    geometries = [mapping(geom) for geom in shapefile.geometry]

    # 创建元数据字典
    meta = {
        "driver": "GTiff",
        "height": data.shape[0],
        "width": data.shape[1],
        "count": 1,
        "dtype": data.dtype,
        "crs": "+proj=latlong",
        "transform": affine,
        "nodata": nodata_value
    }

    # 使用 MemoryFile 创建临时 rasterio 数据集
    with MemoryFile() as mem_file:
        with mem_file.open(**meta) as dataset:
            dataset.write(data, 1)

            # 使用蒙版裁剪栅格数据
            out_data, out_affine = mask(dataset, geometries, crop=True, nodata=nodata_value)
            out_data = out_data[0]  # 提取第一个波段

    return out_data, out_affine


def calc_extent(geo_data, geo_transform, projection):
    # 创建源空间参考
    source_srs = osr.SpatialReference()
    source_srs.ImportFromWkt(projection)
    # 创建目标空间参考（WGS84）
    target_srs = osr.SpatialReference()
    target_srs.ImportFromEPSG(4326)  # WGS84
    # 创建坐标转换对象
    transform = osr.CoordinateTransformation(source_srs, target_srs)
    # 获取数据范围的角点并转换为经纬度
    cols, rows = geo_data.shape[1], geo_data.shape[0]
    lat_lon_corners = [gdal.ApplyGeoTransform(geo_transform, 0, 0),
                       gdal.ApplyGeoTransform(geo_transform, cols, 0),
                       gdal.ApplyGeoTransform(geo_transform, 0, rows),
                       gdal.ApplyGeoTransform(geo_transform, cols, rows)]
    lat_lon_corners = [transform.TransformPoint(x, y) for x, y in lat_lon_corners]
    min_lon, max_lon = min(pt[0] for pt in lat_lon_corners), max(pt[0] for pt in lat_lon_corners)
    min_lat, max_lat = min(pt[1] for pt in lat_lon_corners), max(pt[1] for pt in lat_lon_corners)
    return max_lat, max_lon, min_lat, min_lon


def show_geo_data(geo_data, geo_transform, projection, shp_mask=None, cmap=plt.cm.jet):
    """
    显示地理数据，转换为经纬度坐标系，并应用区域蒙版（如果有）
    :param geo_data: 要显示的地理数据（numpy 数组）
    :param geo_transform: GeoTransform 信息
    :param projection: 投影信息（WKT 格式）
    :param shp_mask: .shp 文件路径，区域蒙版（如果有）
    :param cmap: 色彩映射，默认为 plt.cm.jet
    """
    if shp_mask:
        affine = geo_transform_to_affine(geo_transform)
        geo_data, affine = apply_shapefile_mask(geo_data, affine, shp_mask)
        geo_transform = affine_to_geo_transform(affine)

    max_lat, max_lon, min_lat, min_lon = calc_extent(geo_data, geo_transform, projection)
    extent = [min_lon, max_lon, min_lat, max_lat]

    # 绘制地理数据
    fig, ax = plt.subplots()
    img = ax.imshow(geo_data, extent=extent, cmap=cmap, origin='upper')
    fig.colorbar(img, ax=ax)
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.show()


def export_geo_png(geo_data, geo_transform, projection, png_path, width=None, height=None, full_image=False,
                   shp_mask=None, cmap=plt.cm.jet):
    """
    将 numpy 数组导出为 PNG 图片
    :param geo_data: 要显示的地理数据（numpy 数组）
    :param geo_transform: GeoTransform 信息
    :param projection: 投影信息（WKT 格式）
    :param png_path: 保存 PNG 图片的路径
    :param width: 图片宽度（像素）
    :param height: 图片高度（像素）
    :param full_image: 是否只保存图片而不包含颜色条等（默认 False）
    :param shp_mask: .shp 文件路径，区域蒙版（如果有）
    :param cmap: 色彩映射，默认为 plt.cm.jet
    """
    if shp_mask:
        affine = geo_transform_to_affine(geo_transform)
        geo_data, affine = apply_shapefile_mask(geo_data, affine, shp_mask)
        geo_transform = affine_to_geo_transform(affine)

    max_lat, max_lon, min_lat, min_lon = calc_extent(geo_data, geo_transform, projection)
    extent = [min_lon, max_lon, min_lat, max_lat]

    fig, ax = plt.subplots()
    if width and height:
        fig.set_size_inches(width / PNG_DPI, height / PNG_DPI)  # 以像素为单位指定宽高
    img = ax.imshow(geo_data, extent=extent, cmap=cmap, origin='upper')

    if full_image:
        ax.axis('off')  # 关闭坐标轴
        plt.savefig(png_path, bbox_inches='tight', pad_inches=0, dpi=PNG_DPI)
    else:
        fig.colorbar(img, ax=ax)
        plt.savefig(png_path, bbox_inches='tight', pad_inches=0.1, dpi=PNG_DPI)

    plt.close()
